
import ode
import cairo

class Reference:
	position_x = 0.0
	
	def __init__(self):
		self.step_count = 0
	
	def step(self):
		self.step_count += 1
		self.position_x = self.step_count * 0.25 / 60


class Fitness:
	def __init__(self, simulator):
		self.sim = simulator
		
		self.fitn_center_of_gravity = 0
		self.fitn_length = 0
		self.fitn_energy = 0
		self.fitn_refdiff = 0
		
		self.stepmax = 1200
		self.steps = 0
		self.energy = 0
		self.center_of_gravity_diff = 0
		self.reference_diff = 0
	
	def step(self):
		self.steps += 1
		
		x = self.sim.reference.position_x
		
		#self.energy += abs(self.sim.robot.hip_1.GetJointSpeed() * self.sim.robot.hip_1.GetMotorTorque())
		#self.energy += abs(self.sim.robot.hip_2.GetJointSpeed() * self.sim.robot.hip_2.GetMotorTorque())
		#self.energy += abs(self.sim.robot.knee_1.GetJointSpeed() * self.sim.robot.knee_1.GetMotorTorque())
		#self.energy += abs(self.sim.robot.knee_2.GetJointSpeed() * self.sim.robot.knee_2.GetMotorTorque())
		
		self.center_of_gravity_diff += abs(self.sim.robot.get_center_of_gravity()[0] - x + 0.2)
		
		robopos = (0,0,0) #self.sim.robot.torso.getPosition()
		self.reference_diff += (abs(robopos[0] - x)*1.6) ** 5
		
		self.fitn_center_of_gravity = - self.center_of_gravity_diff / self.steps / 50
		self.fitn_distance = (self.steps ** 1.04) / (self.stepmax)
		self.fitn_energy = - self.energy / self.steps * 1.0
		self.fitn_refdiff = - (self.reference_diff / self.steps) 
		
		
	def get_fitness(self):
		fitness = 2 + \
				  self.fitn_refdiff + \
				  self.fitn_distance + \
				  self.fitn_energy
		if self.fitn_distance > 0.15:
			fitness += (self.fitn_distance-0.15) * 5
			
		if fitness < 0:
			fitness = 0
		return fitness


# Collision callback
def near_callback(args, geom1, geom2):
	"""Callback function for the collide() method.

	This function checks if the given geoms do collide and
	creates contact joints if they do.
	"""

	#if (ode.areConnected(geom1.getBody(), geom2.getBody())):
	#	return


	# Check if the objects do collide
	contacts = ode.collide(geom1, geom2)

	# Create contact joints
	world,contactgroup = args
	for c in contacts:
		c.setBounce(0.2)
		c.setMu(5000)
		j = ode.ContactJoint(world, contactgroup, c)
		j.attach(geom1.getBody(), geom2.getBody())

def matrix2euler(m):
        """Convert from PyODE rotation matrix to euler angles

        Original (in Java) from:
        http://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToEuler/
        """
        import math
        (
                M00, M01, M02,
                M10, M11, M12,
                M20, M21, M22,
        ) = (
                0, 3, 6,
                1, 4, 7,
                2, 5, 8,
        )
        if m[M10] > 0.998:
                y = math.atan2(m[M02], m[M22])
                z = math.pi / 2.0
                x = 0.0
                return (x, y, z)
        if m[M10] < -0.998:
                y = math.atan2(m[M02], m[M22])
                z = -math.pi / 2.0
                x = 0.0
                return (x, y, z)
        y = math.atan2(-m[M02], m[M00])
        x = math.atan2(-m[M12], m[M11])
        z = math.asin(m[M10])
        return (x, y, z)
        
class Simulator:	
	def __init__(self, robot, environment, visualiser = None):
		self.visualiser = visualiser
		self.reference = Reference()
		self.robot = robot
		self.environment = environment
		
		self.world = ode.World()
		self.world.setGravity((0,-9.8,0))
		#self.world.setERP(0.8)
		#self.world.setCFM(1E-3)
		self.space = ode.Space()
		self.floor = ode.GeomPlane(self.space, (0,1,0), 0)
		self.contactgroup = ode.JointGroup()
		
		"""
		world = self.world
		# Create three bodies
		body1 = ode.Body(world)
		M = ode.Mass()
		M.setSphereTotal(200*(10**-9), 50)
		body1.setMass(M)
		body1.setPosition((0,539.5,0))

		body2 = ode.Body(world)
		M = ode.Mass()
		M.setSphereTotal(0.0000002, 50)
		body2.setMass(M)
		body2.setPosition((0,255.5+254/2,0))

		body3 = ode.Body(world)
		M = ode.Mass()
		M.setSphereTotal(0.0000002, 50)
		body3.setMass(M)
		body3.setPosition((0,255.5/2,0))
		
		# Connect body1 with the static environment
		j1 = ode.HingeJoint(world)
		j1.attach(body1, body2)
		j1.setAnchor( (0,255.5+254,0) )
		#Rotate about z axis
		j1.setAxis( (1,0,0) )
		#Set stops for Servo
		j1.setParam(ode.ParamLoStop, -1.6)
		j1.setParam(ode.ParamHiStop,1.6)
		MaxForce = 100

		# Connect body2 with body1
		j2 = ode.HingeJoint(world)
		j2.attach(body2, body3)
		j2.setAnchor( (0,255.5,0) )
		j2.setAxis ((1,0,0))
		j2.setParam(ode.ParamLoStop, -1.6)
		j2.setParam(ode.ParamHiStop,1.6)

		j3 = ode.HingeJoint(world)
		j3.attach(body3, ode.environment)
		j3.setAnchor( (0,0,0) )
		j3.setAxis((1,0,0))
		#j3.setParam(ode.ParamLoStop, -1.6)
		#j3.setParam(ode.ParamHiStop,1.6)
		
		
		self.b1 = body1
		self.b2 = body2
		self.b3 = body3
		self.j1 = j1
		self.j2 = j2
		self.j3 = j3
		"""

		self.robot.add_to_world(self.world, self.space)
		#self.environment.add_to_world(self.world)
		self.fitness = Fitness(self)
	
	def step(self):	
		#self.reference.step()
		#self.robot.step(self.reference)
		#self.fitness.step()
		
		self.space.collide((self.world,self.contactgroup), near_callback)
		self.world.step(1/50.0)
		self.contactgroup.empty()

		if self.visualiser:
			vw = self.visualiser
			
			def buildmatrix(bp, offset):
				P = bp.getPosition()
				P = (P[0]+offset[0], P[1]+offset[1], P[2]+offset[2])
				R = bp.getRotation()
				rot = (R[0], R[3], R[6], 0.,
					R[1], R[4], R[7], 0.,
					R[2], R[5], R[8], 0.,
					P[0], P[1], P[2], 1.0)
				return rot
			
			vw.robot_back.set_matrix(buildmatrix(self.robot.torso, (0, 0,0)))
			vw.robot_hip_joint_l.set_matrix(buildmatrix(self.robot.hip_l, (0,0,0)))
			vw.robot_hip_joint_r.set_matrix(buildmatrix(self.robot.hip_r, (0,0,0)))
			vw.robot_upper_leg_l.set_matrix(buildmatrix(self.robot.upper_leg_l, (0,0,0)))
			vw.robot_upper_leg_r.set_matrix(buildmatrix(self.robot.upper_leg_r, (0,0,0)))
			vw.robot_lower_leg_l.set_matrix(buildmatrix(self.robot.lower_leg_l, (0,0,0)))
			vw.robot_lower_leg_r.set_matrix(buildmatrix(self.robot.lower_leg_r, (0,0,0)))
			
			
			
			"""
			vw.robot_hip_joint_l.set_matrix(buildmatrix(self.b1, (0,0,0)))
			vw.robot_upper_leg_l.set_matrix(buildmatrix(self.b2, (0,0,0)))
			vw.robot_lower_leg_l.set_matrix(buildmatrix(self.b3, (0,0,0)))
			
			
			MaxForce = 0.0000001 #100000000

			#Servo 1
			j1 = self.j1
			j2 = self.j2
			j3 = self.j3
			
			TrueAngle = j3.getAngle()
			DesiredAngle = .8
			Error = TrueAngle - DesiredAngle
			DesiredVelocity = -Error
			j1.setParam(ode.ParamVel, DesiredVelocity)
			j1.setParam(ode.ParamFMax, MaxForce)
			
			TrueAngle = j2.getAngle()
			DesiredAngle = .4
			Error = TrueAngle - DesiredAngle
			DesiredVelocity = -Error
			j2.setParam(ode.ParamVel, DesiredVelocity)
			j2.setParam(ode.ParamFMax, MaxForce)

			#Servo 3
			TrueAngle = j3.getAngle()
			DesiredAngle = +1.2
			Error = TrueAngle - DesiredAngle
			DesiredVelocity = -Error
			j3.setParam(ode.ParamVel, DesiredVelocity)
			j3.setParam(ode.ParamFMax, MaxForce)
			"""

			
			
			vw.draw()
			#print vw.robot_back.position
	
	def get_fitness(self):
		return 1.0 #self.fitness.get_fitness()
